// SimpleConfidentialityToolGUI.java
package com.tidu.demo;

import javax.swing.*;
import javax.swing.filechooser.FileNameExtensionFilter;
import java.awt.*;
import java.io.File;
import java.util.ArrayList;
import java.util.List;

public class SimpleConfidentialityToolGUI extends JFrame {
    private JTextField secretKeyField;
    private JTextField docField;
    private JTextArea logArea;
    private JButton encryptBtn;
    private JButton decryptBtn;
    private JPasswordField passwordField; // 添加密码输入框

    public SimpleConfidentialityToolGUI() {
        setTitle("保密文档脱密加密工具 v1.0 (SM4版本)");
        setDefaultCloseOperation(JFrame.EXIT_ON_CLOSE);
        setSize(700, 550); // 调整窗口大小以适应新增组件
        setLocationRelativeTo(null);
        setResizable(false);

        initUI();
    }

    // 添加新的成员变量
    private JButton batchEncryptBtn;
    private JButton batchDecryptBtn;
    private JTextField folderField;
    private JCheckBox includeSubfoldersCheckBox;
    private JRadioButton useKeyFileRadio; // 使用密钥文件单选按钮
    private JRadioButton usePasswordRadio; // 使用密码单选按钮

    // 修改 initUI 方法，添加批量处理按钮
    private void initUI() {
        JPanel mainPanel = new JPanel(new BorderLayout(10, 10));
        mainPanel.setBorder(BorderFactory.createEmptyBorder(15, 15, 15, 15));

        // 顶部面板 - 文件选择区域
        JPanel topPanel = new JPanel(new GridBagLayout());
        GridBagConstraints gbc = new GridBagConstraints();
        gbc.insets = new Insets(5, 5, 5, 5);
        gbc.anchor = GridBagConstraints.WEST;

        // 密钥选择方式
        gbc.gridx = 0; gbc.gridy = 0;
        topPanel.add(new JLabel("密钥方式:"), gbc);
        gbc.gridx = 1; gbc.gridwidth = 2;

        // 创建单选按钮组
        ButtonGroup keyMethodGroup = new ButtonGroup();
        useKeyFileRadio = new JRadioButton("使用密钥文件", true);
        usePasswordRadio = new JRadioButton("使用密码", false);
        keyMethodGroup.add(useKeyFileRadio);
        keyMethodGroup.add(usePasswordRadio);

        JPanel radioPanel = new JPanel(new FlowLayout(FlowLayout.LEFT));
        radioPanel.add(useKeyFileRadio);
        radioPanel.add(usePasswordRadio);
        topPanel.add(radioPanel, gbc);
        gbc.gridwidth = 1; // 重置gridwidth

        // 密钥文件选择
        gbc.gridx = 0; gbc.gridy = 1;
        topPanel.add(new JLabel("密钥文件:"), gbc);
        gbc.gridx = 1;
        secretKeyField = new JTextField(30);
        topPanel.add(secretKeyField, gbc);
        gbc.gridx = 2;
        JButton secretKeyBtn = new JButton("浏览");
        secretKeyBtn.addActionListener(e -> chooseFile(secretKeyField, "key"));
        topPanel.add(secretKeyBtn, gbc);

        // 密码输入
        gbc.gridx = 0; gbc.gridy = 2;
        topPanel.add(new JLabel("密码:"), gbc);
        gbc.gridx = 1;
        passwordField = new JPasswordField(30);
        topPanel.add(passwordField, gbc);
        gbc.gridx = 2;
        // 添加显示/隐藏密码按钮
        JButton togglePasswordBtn = new JButton("显示");
        togglePasswordBtn.addActionListener(e -> togglePasswordVisibility());
        topPanel.add(togglePasswordBtn, gbc);

        // 单文件选择
        gbc.gridx = 0; gbc.gridy = 3;
        topPanel.add(new JLabel("文档文件:"), gbc);
        gbc.gridx = 1;
        docField = new JTextField(30);
        topPanel.add(docField, gbc);
        gbc.gridx = 2;
        JButton docBtn = new JButton("浏览");
        docBtn.addActionListener(e -> chooseFile(docField, "docx"));
        topPanel.add(docBtn, gbc);

        // 文件夹选择（用于批量处理）
        gbc.gridx = 0; gbc.gridy = 4;
        topPanel.add(new JLabel("文档文件夹:"), gbc);
        gbc.gridx = 1;
        folderField = new JTextField(30);
        topPanel.add(folderField, gbc);
        gbc.gridx = 2;
        JButton folderBtn = new JButton("浏览");
        folderBtn.addActionListener(e -> chooseFolder());
        topPanel.add(folderBtn, gbc);

        // 包含子文件夹选项
        gbc.gridx = 1; gbc.gridy = 5;
        includeSubfoldersCheckBox = new JCheckBox("包含子文件夹", true);
        topPanel.add(includeSubfoldersCheckBox, gbc);

        // 添加事件监听器以控制组件启用状态
        useKeyFileRadio.addActionListener(e -> updateKeyInputComponents());
        usePasswordRadio.addActionListener(e -> updateKeyInputComponents());

        // 初始化组件状态
        updateKeyInputComponents();

        // 按钮面板
        JPanel buttonPanel = new JPanel(new FlowLayout());
        JButton generateBtn = new JButton("生成密钥");
        generateBtn.addActionListener(e -> generateKeys());

        encryptBtn = new JButton("加密文档");
        encryptBtn.addActionListener(e -> encryptDocument());

        decryptBtn = new JButton("解密文档");
        decryptBtn.addActionListener(e -> decryptDocument());

        // 批量处理按钮
        batchEncryptBtn = new JButton("批量加密");
        batchEncryptBtn.addActionListener(e -> batchEncryptDocuments());

        batchDecryptBtn = new JButton("批量解密");
        batchDecryptBtn.addActionListener(e -> batchDecryptDocuments());

        JButton configBtn = new JButton("查看关键词配置");
        configBtn.addActionListener(e -> showConfigInfo());

        buttonPanel.add(generateBtn);
        buttonPanel.add(encryptBtn);
        buttonPanel.add(decryptBtn);
        buttonPanel.add(batchEncryptBtn);
        buttonPanel.add(batchDecryptBtn);
        buttonPanel.add(configBtn);

        // 日志区域
        logArea = new JTextArea(10, 50);
        logArea.setEditable(false);
        logArea.setFont(new Font(Font.MONOSPACED, Font.PLAIN, 12));
        JScrollPane scrollPane = new JScrollPane(logArea);

        mainPanel.add(topPanel, BorderLayout.NORTH);
        mainPanel.add(buttonPanel, BorderLayout.CENTER);
        mainPanel.add(scrollPane, BorderLayout.SOUTH);

        add(mainPanel);
    }

    // 更新密钥输入组件的启用状态
    private void updateKeyInputComponents() {
        boolean useKeyFile = useKeyFileRadio.isSelected();
        secretKeyField.setEnabled(useKeyFile);
        passwordField.setEnabled(!useKeyFile);
    }

    // 切换密码可见性
    private void togglePasswordVisibility() {
        if (passwordField.getEchoChar() == 0) {
            passwordField.setEchoChar('*');
        } else {
            passwordField.setEchoChar((char) 0);
        }
    }

    // 添加文件夹选择方法
    private void chooseFolder() {
        JFileChooser fileChooser = new JFileChooser();
        fileChooser.setFileSelectionMode(JFileChooser.DIRECTORIES_ONLY);
        int result = fileChooser.showOpenDialog(this);
        if (result == JFileChooser.APPROVE_OPTION) {
            folderField.setText(fileChooser.getSelectedFile().getAbsolutePath());
        }
    }

    // 添加获取文件夹中所有docx文件的方法
    private List<String> getDocxFilesFromFolder(String folderPath, boolean includeSubfolders) {
        List<String> files = new ArrayList<>();
        File folder = new File(folderPath);

        if (folder.exists() && folder.isDirectory()) {
            collectDocxFiles(folder, files, includeSubfolders);
        }

        return files;
    }

    private void collectDocxFiles(File folder, List<String> files, boolean includeSubfolders) {
        File[] fileList = folder.listFiles();
        if (fileList != null) {
            for (File file : fileList) {
                if (file.isFile() && (file.getName().toLowerCase().endsWith(".docx") ||
                        file.getName().toLowerCase().endsWith(".doc"))) {
                    files.add(file.getAbsolutePath());
                } else if (file.isDirectory() && includeSubfolders) {
                    collectDocxFiles(file, files, includeSubfolders);
                }
            }
        }
    }

    // 获取密钥字节数组
    private byte[] getSecretKey() throws Exception {
        if (useKeyFileRadio.isSelected()) {
            String secretKeyPath = secretKeyField.getText().trim();
            if (secretKeyPath.isEmpty()) {
                throw new Exception("请选择密钥文件");
            }

            if (!new File(secretKeyPath).exists()) {
                throw new Exception("密钥文件不存在: " + secretKeyPath);
            }

            return ConfidentialityProcessor.loadSecretKey(secretKeyPath);
        } else {
            String password = new String(passwordField.getPassword()).trim();
            if (password.isEmpty()) {
                throw new Exception("请输入密码");
            }

            // 使用固定盐值（在实际应用中应该使用随机盐值并存储）
            byte[] salt = "ConfidentialToolSalt".getBytes();
            return ConfidentialityProcessor.deriveKeyFromPassword(password, salt);
        }
    }

    // 批量加密文档
    private void batchEncryptDocuments() {
        String folderPath = folderField.getText().trim();

        if (folderPath.isEmpty()) {
            JOptionPane.showMessageDialog(this, "请选择文档文件夹", "提示", JOptionPane.WARNING_MESSAGE);
            return;
        }

        File folder = new File(folderPath);
        if (!folder.exists() || !folder.isDirectory()) {
            JOptionPane.showMessageDialog(this, "文件夹不存在: " + folderPath, "错误", JOptionPane.ERROR_MESSAGE);
            return;
        }

        List<String> docxFiles = getDocxFilesFromFolder(folderPath, includeSubfoldersCheckBox.isSelected());
        if (docxFiles.isEmpty()) {
            JOptionPane.showMessageDialog(this, "文件夹中没有找到.docx文件", "提示", JOptionPane.WARNING_MESSAGE);
            return;
        }

        batchEncryptBtn.setEnabled(false);
        new Thread(() -> {
            try {
                logMessage("开始批量加密 " + docxFiles.size() + " 个文档...");
                byte[] secretKey = getSecretKey();

                ConfidentialityProcessor.ProgressCallback callback = new ConfidentialityProcessor.ProgressCallback() {
                    @Override
                    public void onProgress(int current, int total, String message) {
                        SwingUtilities.invokeLater(() -> {
                            logMessage(message);
                        });
                    }

                    @Override
                    public void onSuccess(int index, String filePath) {
                        SwingUtilities.invokeLater(() -> {
                            logMessage("成功加密: " + filePath);
                        });
                    }

                    @Override
                    public void onError(int index, String filePath, Exception error) {
                        SwingUtilities.invokeLater(() -> {
                            logMessage("加密失败: " + filePath + " - " + error.getMessage());
                        });
                    }
                };

                ConfidentialityProcessor.batchEncryptConfidentialContent(docxFiles, secretKey, callback);
                logMessage("批量加密完成!");

                SwingUtilities.invokeLater(() -> {
                    JOptionPane.showMessageDialog(this,
                            "批量加密完成!\n成功处理 " + docxFiles.size() + " 个文档",
                            "成功", JOptionPane.INFORMATION_MESSAGE);
                });
            } catch (Exception ex) {
                logMessage("批量加密过程中出错: " + ex.getMessage());
                SwingUtilities.invokeLater(() -> {
                    JOptionPane.showMessageDialog(this,
                            "批量加密过程中出错: " + ex.getMessage(),
                            "错误", JOptionPane.ERROR_MESSAGE);
                });
            } finally {
                SwingUtilities.invokeLater(() -> batchEncryptBtn.setEnabled(true));
            }
        }).start();
    }

    // 批量解密文档
    private void batchDecryptDocuments() {
        String folderPath = folderField.getText().trim();

        if (folderPath.isEmpty()) {
            JOptionPane.showMessageDialog(this, "请选择文档文件夹", "提示", JOptionPane.WARNING_MESSAGE);
            return;
        }

        File folder = new File(folderPath);
        if (!folder.exists() || !folder.isDirectory()) {
            JOptionPane.showMessageDialog(this, "文件夹不存在: " + folderPath, "错误", JOptionPane.ERROR_MESSAGE);
            return;
        }

        List<String> docxFiles = getDocxFilesFromFolder(folderPath, includeSubfoldersCheckBox.isSelected());
        if (docxFiles.isEmpty()) {
            JOptionPane.showMessageDialog(this, "文件夹中没有找到.docx文件", "提示", JOptionPane.WARNING_MESSAGE);
            return;
        }

        batchDecryptBtn.setEnabled(false);
        new Thread(() -> {
            try {
                logMessage("开始批量解密 " + docxFiles.size() + " 个文档...");
                byte[] secretKey = getSecretKey();

                ConfidentialityProcessor.ProgressCallback callback = new ConfidentialityProcessor.ProgressCallback() {
                    @Override
                    public void onProgress(int current, int total, String message) {
                        SwingUtilities.invokeLater(() -> {
                            logMessage(message);
                        });
                    }

                    @Override
                    public void onSuccess(int index, String filePath) {
                        SwingUtilities.invokeLater(() -> {
                            logMessage("成功解密: " + filePath);
                        });
                    }

                    @Override
                    public void onError(int index, String filePath, Exception error) {
                        SwingUtilities.invokeLater(() -> {
                            logMessage("解密失败: " + filePath + " - " + error.getMessage());
                        });
                    }
                };

                ConfidentialityProcessor.batchDecryptConfidentialContent(docxFiles, secretKey, callback);
                logMessage("批量解密完成!");

                SwingUtilities.invokeLater(() -> {
                    JOptionPane.showMessageDialog(this,
                            "批量解密完成!\n成功处理 " + docxFiles.size() + " 个文档",
                            "成功", JOptionPane.INFORMATION_MESSAGE);
                });
            } catch (Exception ex) {
                logMessage("批量解密过程中出错: " + ex.getMessage());
                SwingUtilities.invokeLater(() -> {
                    JOptionPane.showMessageDialog(this,
                            "批量解密过程中出错: " + ex.getMessage(),
                            "错误", JOptionPane.ERROR_MESSAGE);
                });
            } finally {
                SwingUtilities.invokeLater(() -> batchDecryptBtn.setEnabled(true));
            }
        }).start();
    }

    private void chooseFile(JTextField field, String fileType) {
        JFileChooser fileChooser = new JFileChooser();
        if ("key".equals(fileType)) {
            fileChooser.setFileFilter(new FileNameExtensionFilter("密钥文件 (*.key)", "key"));
        } else if ("doc".equals(fileType)) {
            fileChooser.setFileFilter(new FileNameExtensionFilter("Word文档 (*.doc, *.docx)", "doc", "docx"));
        }

        int result = fileChooser.showOpenDialog(this);
        if (result == JFileChooser.APPROVE_OPTION) {
            field.setText(fileChooser.getSelectedFile().getAbsolutePath());
        }
    }

    private void logMessage(String message) {
        logArea.append("[" + new java.util.Date() + "] " + message + "\n");
        logArea.setCaretPosition(logArea.getDocument().getLength());
    }

    private void generateKeys() {
        JFileChooser fileChooser = new JFileChooser();
        fileChooser.setDialogTitle("选择密钥保存位置");
        fileChooser.setSelectedFile(new File("sm4.key"));

        int result = fileChooser.showSaveDialog(this);
        if (result == JFileChooser.APPROVE_OPTION) {
            try {
                String secretKeyPath = fileChooser.getSelectedFile().getAbsolutePath();
                if (!secretKeyPath.endsWith(".key")) {
                    secretKeyPath += ".key";
                }

                logMessage("正在生成SM4密钥...");
                byte[] secretKey = ConfidentialityProcessor.generateSM4Key();
                ConfidentialityProcessor.saveSecretKey(secretKey, secretKeyPath);

                logMessage("SM4密钥生成成功!");
                logMessage("密钥保存到: " + secretKeyPath);

                JOptionPane.showMessageDialog(this,
                        "SM4密钥生成成功!\n密钥保存到: " + secretKeyPath,
                        "成功", JOptionPane.INFORMATION_MESSAGE);
            } catch (Exception ex) {
                logMessage("生成SM4密钥时出错: " + ex.getMessage());
                JOptionPane.showMessageDialog(this,
                        "生成SM4密钥时出错: " + ex.getMessage(),
                        "错误", JOptionPane.ERROR_MESSAGE);
            }
        }
    }

    private void encryptDocument() {
        String srcPath = docField.getText().trim();

        if (srcPath.isEmpty()) {
            JOptionPane.showMessageDialog(this, "请选择文档文件", "提示", JOptionPane.WARNING_MESSAGE);
            return;
        }

        if (!new File(srcPath).exists()) {
            JOptionPane.showMessageDialog(this, "文档文件不存在: " + srcPath, "错误", JOptionPane.ERROR_MESSAGE);
            return;
        }

        encryptBtn.setEnabled(false);
        new Thread(() -> {
            try {
                logMessage("开始加密文档: " + srcPath);
                byte[] secretKey = getSecretKey();
                ConfidentialityProcessor.encryptConfidentialContent(srcPath, secretKey);

                logMessage("文档加密完成!");

                SwingUtilities.invokeLater(() -> {
                    JOptionPane.showMessageDialog(this, "文档加密完成!\n",
                            "成功", JOptionPane.INFORMATION_MESSAGE);
                });
            } catch (Exception ex) {
                logMessage("加密过程中出错: " + ex.getMessage());
                SwingUtilities.invokeLater(() -> {
                    JOptionPane.showMessageDialog(this,
                            "加密过程中出错: " + ex.getMessage(),
                            "错误", JOptionPane.ERROR_MESSAGE);
                });
            } finally {
                SwingUtilities.invokeLater(() -> encryptBtn.setEnabled(true));
            }
        }).start();
    }

    private void decryptDocument() {
        String srcPath = docField.getText().trim();

        if (srcPath.isEmpty()) {
            JOptionPane.showMessageDialog(this, "请选择文档文件", "提示", JOptionPane.WARNING_MESSAGE);
            return;
        }

        if (!new File(srcPath).exists()) {
            JOptionPane.showMessageDialog(this, "文档文件不存在: " + srcPath, "错误", JOptionPane.ERROR_MESSAGE);
            return;
        }

        decryptBtn.setEnabled(false);
        new Thread(() -> {
            try {
                logMessage("开始解密文档: " + srcPath);
                byte[] secretKey = getSecretKey();

                ConfidentialityProcessor.decryptConfidentialContent(srcPath, secretKey);

                logMessage("文档解密完成!");

                SwingUtilities.invokeLater(() -> {
                    JOptionPane.showMessageDialog(this, "文档解密完成!", "成功", JOptionPane.INFORMATION_MESSAGE);
                });
            } catch (Exception ex) {
                logMessage("解密过程中出错: " + ex.getMessage());
                SwingUtilities.invokeLater(() -> {
                    JOptionPane.showMessageDialog(this,
                            "解密过程中出错: " + ex.getMessage(),
                            "错误", JOptionPane.ERROR_MESSAGE);
                });
            } finally {
                SwingUtilities.invokeLater(() -> decryptBtn.setEnabled(true));
            }
        }).start();
    }

    private void showConfigInfo() {
        try {
            List<String> keywords = ConfigManager.loadCustomKeywords();
            StringBuilder info = new StringBuilder();
            info.append("当前配置的涉密关键词:\n\n");
            for (int i = 0; i < keywords.size(); i++) {
                info.append((i + 1)).append(". ").append(keywords.get(i)).append("\n");
            }
            info.append("\n如需修改，请编辑程序目录下的 confidential_keywords.txt 文件");

            JTextArea textArea = new JTextArea(info.toString());
            textArea.setEditable(false);
            textArea.setRows(15);
            textArea.setColumns(40);
            // 确保使用正确的字体显示中文
            textArea.setFont(new Font("Microsoft YaHei", Font.PLAIN, 12));
            JScrollPane scrollPane = new JScrollPane(textArea);

            JOptionPane.showMessageDialog(this, scrollPane, "关键词配置信息", JOptionPane.INFORMATION_MESSAGE);
        } catch (Exception ex) {
            JOptionPane.showMessageDialog(this, "读取配置信息出错: " + ex.getMessage(),
                    "错误", JOptionPane.ERROR_MESSAGE);
        }
    }

    public static void main(String[] args) {
        SwingUtilities.invokeLater(() -> {
            try {
                UIManager.setLookAndFeel(UIManager.getSystemLookAndFeelClassName());
            } catch (Exception e) {
                e.printStackTrace();
            }
            new SimpleConfidentialityToolGUI().setVisible(true);
        });
    }
}
